import cv2
import numpy as np
import odl
import scipy.io
import torch
from odl.contrib import torch as odl_torch


## 640geo
class initialization:
    def __init__(self):
        self.param = {}
        self.reso = 512 / 416 * 0.03

        # image
        # self.param['nx_h'] = 416
        # self.param['ny_h'] = 416
        self.param['nx_h'] = 364  # modified by Gao
        self.param['ny_h'] = 364
        self.param['sx'] = self.param['nx_h'] * self.reso
        self.param['sy'] = self.param['ny_h'] * self.reso

        ## view
        self.param['startangle'] = 0
        self.param['endangle'] = 2 * np.pi

        # self.param['nProj'] = 640
        self.param['nProj'] = 720

        ## detector
        self.param['su'] = 2 * np.sqrt(self.param['sx'] ** 2 + self.param['sy'] ** 2)
        self.param['nu_h'] = 641
        # self.param['dde'] = 1075 * self.reso
        # self.param['dso'] = 1075 * self.reso
        self.param['dde'] = 1  # diameter of the detector?
        self.param['dso'] = 1  # diameter of the source?

        self.param['u_water'] = 0.192


def build_gemotry(param):
    """

    """
    reco_space_h = odl.uniform_discr(
        min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
        max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
        dtype='float32')
    # startangle 0 endangle 2Pi nProj 720
    angle_partition = odl.uniform_partition(param.param['startangle'], param.param['endangle'],
                                            param.param['nProj'])
    # su 38.0140605566 nu_h 641
    detector_partition_h = odl.uniform_partition(-(param.param['su'] / 2.0), (param.param['su'] / 2.0),
                                                 param.param['nu_h'])
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          src_radius=param.param['dso'],  # Radius of the source circle
                                          det_radius=param.param['dde'])  # Radius of the detector circle

    ray_trafo_hh = odl.tomo.RayTransform(reco_space_h, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


def build_gemotry_own():
    """
    目前锥束加入到网络中都比较难，需要的显存过多。
    因此测试使用deeplesion数据，使用扇形束即可，但是参数和InDuDoNet使用的可能不太一样。
    """
    # reco_space_h = odl.uniform_discr(
    #     min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
    #     max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
    #     dtype='float32')
    # Todo 还没搞懂这个参数是干啥用的，先用官方代码中的示例
    reco_space = odl.uniform_discr(
        min_pt=[-20, -20], max_pt=[20, 20], shape=[256, 256], dtype='float32')
    # startangle 0 endangle 2Pi nProj 720  # 320来自于ADN中制作deeplesion数据时使用的参数（光源旋转次数）
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 180)
    # su 38.0140605566 nu_h 641  0.171来自于ADN中制作deeplesion数据时matlab函数para2fan源码中gammaRad变量
    detector_partition_h = odl.uniform_partition(-60, 60, 240)
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          src_radius=100,  # Radius of the source circle
                                          det_radius=100)  # Radius of the detector circle

    ray_trafo_hh = odl.tomo.RayTransform(reco_space, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


def build_gemotry_rat_femur():
    """
    目前锥束加入到网络中都比较难，需要的显存过多。
    因此测试使用deeplesion数据，使用扇形束即可，但是参数和InDuDoNet使用的可能不太一样。
    """
    # reco_space_h = odl.uniform_discr(
    #     min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
    #     max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
    #     dtype='float32')
    # Todo 还没搞懂这个参数是干啥用的，先用官方代码中的示例
    reco_space = odl.uniform_discr(
        min_pt=[-20, -20], max_pt=[20, 20], shape=[450, 450], dtype='float32')
    # startangle 0 endangle 2Pi nProj 720  # 320来自于ADN中制作deeplesion数据时使用的参数（光源旋转次数）
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 180)
    # su 38.0140605566 nu_h 641  0.171来自于ADN中制作deeplesion数据时matlab函数para2fan源码中gammaRad变量
    detector_partition_h = odl.uniform_partition(-60, 60, 240)
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          src_radius=100,  # Radius of the source circle
                                          det_radius=100)  # Radius of the detector circle

    ray_trafo_hh = odl.tomo.RayTransform(reco_space, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


def build_gemotry_deeplesion():
    """
    目前锥束加入到网络中都比较难，需要的显存过多。
    因此测试使用deeplesion数据，使用扇形束即可，但是参数和InDuDoNet使用的可能不太一样。
    """
    # reco_space_h = odl.uniform_discr(
    #     min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
    #     max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
    #     dtype='float32')
    # Todo 还没搞懂这个参数是干啥用的，先用官方代码中的示例
    reso = 0.05
    reco_space = odl.uniform_discr(
        min_pt=[-128 * reso, -128 * reso], max_pt=[128 * reso, 128 * reso], shape=[256, 256], dtype='float32')
    # startangle 0 endangle 2Pi nProj 720  # 320来自于ADN中制作deeplesion数据时使用的参数（光源旋转次数）
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 320)
    # su 38.0140605566 nu_h 641  0.171来自于ADN中制作deeplesion数据时matlab函数para2fan源码中gammaRad变量
    detector_partition_h = odl.uniform_partition(-184 * reso, 184 * reso, shape=369)
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          # src_to_det_init=(0, 1075),
                                          src_radius=1075 * reso,  # Radius of the source circle
                                          det_radius=0,
                                          )  # Radius of the detector circle
    print(geometry_h.src_position(0))
    ray_trafo_hh = odl.tomo.RayTransform(reco_space, geometry_h, impl='astra_cuda')
    return ray_trafo_hh


def generate_deeplesion_odl_operator():
    """
    实验发现，如果没有滤波的话，重建出来的图像的值域会比原始图的大很多
    """
    fp = build_gemotry_deeplesion()  # fp是一个RayTransform类
    # fp = build_gemotry_own()
    op_module_fp = odl_torch.OperatorModule(fp)  # OperatorModule类用来封装
    op_module_pT = odl_torch.OperatorModule(fp.adjoint)
    return op_module_fp, op_module_pT


def generate_deeplesion_odl_operator_old():
    """
    实验发现，如果没有滤波的话，重建出来的图像的值域会比原始图的大很多
    """
    # fp = build_gemotry_deeplesion()  # fp是一个RayTransform类
    fp = build_gemotry_own()
    op_module_fp = odl_torch.OperatorModule(fp)  # OperatorModule类用来封装
    op_module_pT = odl_torch.OperatorModule(fp.adjoint)
    return op_module_fp, op_module_pT


def generate_rat_femur_odl_operator():
    fp = build_gemotry_rat_femur()  # fp是一个RayTransform类
    op_module_fp = odl_torch.OperatorModule(fp)  # OperatorModule类用来封装
    op_module_pT = odl_torch.OperatorModule(fp.adjoint)
    return op_module_fp, op_module_pT


def test_odl_fbp_function():  # --- Set up geometry of the problem --- #

    # Reconstruction space: discretized functions on the cube
    # [-20, 20]^2 with 300 samples per dimension.
    reco_space = odl.uniform_discr(
        min_pt=[-20, -20], max_pt=[20, 20], shape=[300, 300],
        dtype='float32')

    # Make a circular cone beam geometry with flat detector
    # Angles: uniformly spaced, n = 360, min = 0, max = 2 * pi
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 360)
    # Detector: uniformly sampled, n = 512, min = -60, max = 60
    detector_partition = odl.uniform_partition(-60, 60, 512)
    # Geometry with large fan angle
    geometry = odl.tomo.FanBeamGeometry(
        angle_partition, detector_partition, src_radius=40, det_radius=40)

    # --- Create Filtered Back-projection (FBP) operator --- #

    # Ray transform (= forward projection).
    ray_trafo = odl.tomo.RayTransform(reco_space, geometry)

    # Create FBP operator using utility function
    # We select a Hann filter, and only use the lowest 80% of frequencies to avoid
    # high frequency noise.
    fbp = odl.tomo.fbp_op(ray_trafo, filter_type='Hann', frequency_scaling=0.8)

    # --- Show some examples --- #

    # Create a discrete Shepp-Logan phantom (modified version)
    phantom = odl.phantom.shepp_logan(reco_space, modified=True)

    # Create projection data by calling the ray transform on the phantom
    proj_data = ray_trafo(phantom)

    # Calculate filtered back-projection of data
    fbp_reconstruction = fbp(proj_data)

    # Shows a slice of the phantom, projections, and reconstruction
    phantom.show(title='Phantom')
    phantom_ndarray = phantom.data
    proj_data.show(title='Projection Data (Sinogram)')
    proj_data_ndarray = proj_data.data
    fbp_reconstruction.show(title='Filtered Back-projection')
    fbp_reconstruction_ndarray = fbp_reconstruction.data
    (phantom - fbp_reconstruction).show(title='Error', force_show=True)
    print("good")


def test_odl_function():
    """Example using the ray transform with fan beam geometry."""

    # Reconstruction space: discretized functions on the rectangle
    # [-20, 20]^2 with 300 samples per dimension.
    reco_space = odl.uniform_discr(
        min_pt=[-20, -20], max_pt=[20, 20], shape=[300, 300], dtype='float32')

    # Make a fan beam geometry with flat detector
    # Angles: uniformly spaced, n = 360, min = 0, max = 2 * pi
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 360)
    # Detector: uniformly sampled, n = 512, min = -30, max = 30
    detector_partition = odl.uniform_partition(-30, 30, 512)
    geometry = odl.tomo.FanBeamGeometry(angle_partition, detector_partition,
                                        src_radius=1000, det_radius=100)

    # Ray transform (= forward projection).
    ray_trafo = odl.tomo.RayTransform(reco_space, geometry)

    # Create a discrete Shepp-Logan phantom (modified version)
    phantom = odl.phantom.shepp_logan(reco_space, modified=True)

    # Create projection data by calling the ray transform on the phantom
    proj_data = ray_trafo(phantom)

    # Back-projection can be done by simply calling the adjoint operator on the
    # projection data (or any element in the projection space).
    backproj = ray_trafo.adjoint(proj_data)

    # op_module_fp = odl_torch.OperatorModule(ray_trafo)  # OperatorModule类用来封装
    # op_module_pT = odl_torch.OperatorModule(ray_trafo.adjoint)
    # my_proj_data = op_module_fp(phantom)
    # my_recon = op_module_pT(my_proj_data)
    # Shows a slice of the phantom, projections, and reconstruction
    phantom.show(title='Phantom')
    phantom_ndarray = phantom.data
    proj_data.show(title='Projection Data (sinogram)')
    proj_data_ndarray = proj_data.data
    backproj.show(title='Back-projection', force_show=True)
    backproj_ndarray = backproj.data
    print("good")


if __name__ == '__main__':
    """matlab中，光源的旋转方向为逆时针
    """
    # test_odl_function()
    op_module_fp, op_module_pT = generate_deeplesion_odl_operator()
    op_module_fp_old, op_module_pT_old = generate_deeplesion_odl_operator_old()
    test_file_path = r'E:\data_transform_buffer\metal_artifact\easy_dataset\test\A\000365_03_01_473_2.png'
    mat_file_path = r"F:\metal_artifact_data\deeplesion\original_data\Images_png\Images_png_01\train\000001_01_01\104\ma_CT_3.mat"
    mat_file = scipy.io.loadmat(mat_file_path)['ma_CT']
    # mat_file = (mat_file + 1) / 5 * 255
    mat_file = np.clip(mat_file, 0, 1)
    test_png_tensor = torch.tensor(mat_file).unsqueeze(0).unsqueeze(0)
    test_project = op_module_fp(test_png_tensor)
    test_recon = op_module_pT(test_project)
    test_project = test_project.numpy().squeeze()
    test_recon_show = test_recon.numpy().squeeze()

    test_project_old = op_module_fp_old(test_png_tensor)
    test_recon_old = op_module_pT_old(test_project_old)
    test_project_old = test_project_old.numpy().squeeze()
    test_recon_show_old = test_recon_old.numpy().squeeze()

    cv2.imshow("test_project", (test_project / 70 * 255).astype(np.uint8))
    cv2.imshow("test_recon_show", test_recon_show / 30000)
    cv2.imshow("test_project_old", (test_project_old / 12 * 255).astype(np.uint8))
    cv2.imshow("test_recon_show_old", test_recon_show_old / 50)
    cv2.waitKey()
    cv2.destroyAllWindows()

    test_png = cv2.imread(test_file_path, 0) / 255.
    # test_png = np.transpose(test_png,)

    cv2.imshow("test_png", test_png)
    scipy.io.savemat("test_project.mat", {"test_project": test_project})
    cv2.imshow("test_project", (test_project / 123 * 255).astype(np.uint8))
    cv2.imshow("test_recon_show", test_recon_show / 30000)
    cv2.waitKey()
    cv2.destroyAllWindows()
